package Test.rbTree;

/**
 * 红黑树手写实现
 * <p>
 * ①创建RBTree，定义颜色
 * ②创建RBNode
 * ③辅助方法定义：parentOf(node)，isRed(node)，setRed(node)，setBlack(node)，inOrderPrint(RBNode tree)
 * ④左旋方法定义：leftRotate(node)
 * ⑤右旋方法定义：rightRotate(node)
 * ⑥公开插入接口方法定义：insert(K key, V value);
 * ⑦内部插入接口方法定义：insert(RBNode node);
 * ⑧修正插入导致红黑树失衡的方法定义：insertFIxUp(RBNode node);
 * ⑨测试红黑树正确性
 *
 *
 */
public class RBTree<k extends Comparable<k>, v> {

    private static final boolean RED = true;
    private static final boolean BLACK = false;

    /**
     * 根节点
     **/
    private RBNode root;

    public RBNode getRoot() {
        return root;
    }

    /**
     * 插入之后修复红黑树，保证红黑树结构完整的方法
     *
     * 情景1：红黑树只有一个节点
     *      插入的节点是红色的，这个时候需要把这个节点染色
     * 情景2：插入节点的key已经存在
     *      无法进入，不需要处理
     * 情景3：插入节点的父节点为黑色
     *      由于没有改变红黑树的结构，所以不需要处理
     * 情景4：插入节点的父节点为红色
     *      破坏了红黑树的结构，需要根据具体的情况进行处理
     *      4.1 父节点和叔叔节点同时存在且为红色
     *      下面的情况，叔叔节点要么不存在，要么是黑色
     *      4.2 父节点在爷爷节点的左边，要插入节点在父节点的左边
     *      4.3 父节点在爷爷节点的左边，要插入节点在父节点的右边
     *      4.4 父节点在爷爷节点的右边，要插入节点在父节点的右边
     *      4.5 父节点在爷爷节点的右边，要插入节点在父节点的左边
     *
     */
    private void insertFixUp(RBNode node) {
        //每次设置根节点的颜色，处理情景1
        this.root.color = BLACK;

        RBNode parent = parentOf(node);
        RBNode grand = parentOf(parent);
        //情景4
        if (parent != null && isRed(parent)){
            RBNode uncle = null;

            //如果父节点为红色，那么一定存在爷爷节点，并且爷爷节点一定为黑色
            if (parent == grand.left){
                uncle = grand.right;
                if (uncle != null  && isRed(uncle)){
                    //情景4.1 父节点和叔叔节点同时存在且为红色
                    //把 父节点叔叔节点爷爷节点全部变色
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(grand);
                    insertFixUp(grand);
                    return;
                }
                if (uncle == null || uncle.color == BLACK){
                    if (node == parent.left){
                        //情景4.2 父节点在爷爷节点的左边，要插入节点在父节点的左边
                        setBlack(parent);
                        setRed(grand);
                        rightRotate(grand);
                        return;
                    }else {
                        //情景4.3 父节点在爷爷节点的左边，要插入节点在父节点的右边
                        leftRotate(parent);
                        //直接以parent节点进入下一轮的处理
                        insertFixUp(parent);
                        return;
                    }
                }
            }else {
                uncle = grand.left;
                if (uncle != null  && isRed(uncle)){
                    //情景4.1 父节点和叔叔节点同时存在且为红色
                    //把 父节点叔叔节点爷爷节点全部变色
                    setBlack(parent);
                    setBlack(uncle);
                    setRed(grand);
                    insertFixUp(grand);
                    return;
                }

                if (uncle == null || uncle.color == BLACK){

                    if (node == parent.right){
                        //情景4.4 父节点在爷爷节点的右边，要插入节点在父节点的右边
                        setBlack(parent);
                        setRed(grand);
                        leftRotate(grand);
                        return;
                    }

                    if (node == parent.left){
                        //情景4.5 父节点在爷爷节点的右边，要插入节点在父节点的左边
                        rightRotate(parent);
                        insertFixUp(parent);
                        return;
                    }
                }
            }
        }
        //父节点不存在，或者父节点是黑色，不需要处理自动跳出
    }


    /**
     * 对外提供的插入方法
     * @param key key
     * @param value value
     * @return 返回值,如果插入成功返回null 如果插入失败，即发生了替换的操作则返回旧值
     */
    public v insert(k key,v value){
        RBNode node = new RBNode();
        node.key = key;
        node.value = value;
        node.color = RED;
        return insert(node);
    }

    /**
     * 核心插入方法
     * @param node 要插入的节点
     * @return 返回值,如果插入成功返回null 如果插入失败，即发生了替换的操作则返回旧值
     */
    private v insert(RBNode node){
        //1.寻找父节点
        RBNode x = this.root;
        RBNode parent = null;

        while (x!= null) {
            parent = x;
            //1. cmp=0 key相同，找到了相同的key，需要执行替换操作
            //2. cmp>0 需要向左子树搜索
            //3. cmp<0 需要向右子树搜索
            int cmp = x.key.compareTo(node.key);
            if (cmp == 0){
                v old = (v)x.value;
                x.value = node.value;
                return old;
            }else if (cmp > 0){
                x = x.left;
            }else {
                x = x.right;
            }
        }

        if (parent != null) {
            //找到了父节点，通过key的大小，判断node是父节点的左还是右
            node.parent = parent;
            if (node.key.compareTo(parent.key) > 0) {
                parent.right = node;
            } else if (node.key.compareTo(parent.key) < 0) {
                parent.left = node;
            }
        }else {
            //没有找到父节点，换句话说就是整个红黑树一个节点都没有
            this.root = node;
        }

        //调用修复红黑树的方法
        insertFixUp(node);
        return null;
    }

    /**
     * 右旋方法
     * 右旋示意图：右旋y节点
     *
     *    p                       p
     *    |                       |
     *    y                       x
     *   / \          ---->      / \
     *  x   ry                  lx  y
     * / \                         / \
     *lx  ly                      ly  ry
     *
     * 右旋做了那些事情：
     * 1.把x的右子节点变成y的左子节点，把x的右子节点指向y
     * 2.当y的父节点存在的时候，把x的父节点只想p，把p的子节点指向x
     * 3.y的父节点指向x，把x的右子节点指向y
     *
     * @param node 要旋转的节点
     */
    private void rightRotate(RBNode node) {
        RBNode leftNode = node.left;
        //子节点的处理
        node.left = leftNode.right;
        if (leftNode.right != null) {
            leftNode.right.parent = node;
        }
        //父节点的处理
        if (node.parent != null) {
            if (node.parent.left == node) {
                node.parent.left = leftNode;
            } else {
                node.parent.right = leftNode;
            }
            leftNode.parent = node.parent;
        } else {
            this.root = leftNode;
            this.root.parent = null;
        }
        //收尾工作 y的父节点变成x  x的右子节点变成y
        node.parent = leftNode;
        leftNode.right = node;
    }


    /**
     * 左旋方法
     * 左旋示意图：左旋x节点
     *
     *    p                   p
     *    |                   |
     *    x                   y
     *   / \         ---->   / \
     *  lx  y               x   ry
     *     / \             / \
     *    ly  ry          lx  ly
     *
     * 左旋做了那些事情：
     * 1.把y的左子节点的父节点变成x，把x的右子节点变成y的左子节点
     * 2.当x的父节点存在的时候，把y的父节变成p，把p指向y
     * 3.把x的父节点变成y，把y的左子节点变成x
     *
     * @param node 要旋转的节点
     */
    private void leftRotate(RBNode node) {
        //右子节点
        RBNode rightNode = node.right;
        //把x的右子节点变成y的左子节点
        node.right = rightNode.left;
        //把y的左子节点的父节点变成x
        if (rightNode.left != null) {
            rightNode.left.parent = node;
        }
        //父节点处理
        if (node.parent != null) {
            //新的节点只想被旋转节点的父节点
            rightNode.parent = node.parent;
            //父节点指向新的节点
            if (node.parent.left == node) {
                node.parent.left = rightNode;
            } else {
                node.parent.right = rightNode;
            }
        } else {
            //父节点是根节点的情况
            this.root = rightNode;
            this.root.parent = null;
        }

        node.parent = rightNode;
        rightNode.left = node;
    }

    /**
     * 中序打印红黑树
     */
    public void inOrderPrint() {
        inOrderPrint(root);
    }

    /**
     * 中序打印红黑树
     *
     * @param node 节点
     */
    private void inOrderPrint(RBNode node) {
        if (node != null) {
            inOrderPrint(node.left);
            System.out.println("key =>" + node.key + ",value =>" + node.value);
            inOrderPrint(node.right);
        }
    }

    /**
     * 设置节点为黑色
     *
     * @param node 节点
     */
    private void setBlack(RBNode node) {
        if (node != null) {
            node.color = BLACK;
        }
    }

    /**
     * 设置节点为红色
     *
     * @param node 节点
     */
    private void setRed(RBNode node) {
        if (node != null) {
            node.color = RED;
        }
    }

    /**
     * 判断是否是红色节点
     * 由于 null 也被判断为黑色节点，所以node为空返回 false
     *
     * @param node 节点
     * @return 节点类型
     */
    private boolean isRed(RBNode node) {
        if (node != null) {
            return node.color;
        }
        return false;
    }

    /**
     * 获取父节点
     *
     * @param node 节点
     * @return 父节点
     */
    private RBNode parentOf(RBNode node) {
        if (node != null) {
            return node.parent;
        }
        return null;
    }

    /**
     * 树节点内部类
     *
     * @param <k> key
     * @param <v> value
     */
    static class RBNode<k extends Comparable<k>, v> {
        private RBNode<k, v> parent;
        private RBNode<k, v> left;
        private RBNode<k, v> right;
        private k key;
        private v value;
        private boolean color;

        public RBNode() {
        }

        public RBNode(RBNode<k, v> parent, RBNode<k, v> left, RBNode<k, v> right, k key, v value, boolean color) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.key = key;
            this.value = value;
            this.color = color;
        }

        public RBNode<k, v> getParent() {
            return parent;
        }

        public void setParent(RBNode<k, v> parent) {
            this.parent = parent;
        }

        public RBNode<k, v> getLeft() {
            return left;
        }

        public void setLeft(RBNode<k, v> left) {
            this.left = left;
        }

        public RBNode<k, v> getRight() {
            return right;
        }

        public void setRight(RBNode<k, v> right) {
            this.right = right;
        }

        public k getKey() {
            return key;
        }

        public void setKey(k key) {
            this.key = key;
        }

        public v getValue() {
            return value;
        }

        public void setValue(v value) {
            this.value = value;
        }

        public boolean isColor() {
            return color;
        }

        public void setColor(boolean color) {
            this.color = color;
        }
    }

}